import torch.nn as nn 
import torch 

class Discriminator(nn.Module):

    def __init__(self, device="cpu"):
        super(Discriminator, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 128, (3, 3), stride=2, padding=(1, 1)),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.4),
            
            nn.Conv2d(128, 128, (3, 3), stride=2, padding=(1, 1)),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.4),
        )
        self.output = nn.Linear(128 * 7 * 7, 1)
        self.sigmoid = nn.Sigmoid()
        self.device = device

    def forward(self, x, labels=None):
        x = x.to(self.device)
        ## x: (batch_size, 1, 28, 28)
        # x = x.view(x.shape[0], 1, x.shape[-1], x.shape[-1])
        x = self.net(x)
        x = x.view(-1, 128 * 7 * 7)
        x = self.output(x)
        x = self.sigmoid(x)
        
        return x.view(-1)

class Generator(nn.Module):
    
    def __init__(self, latent_dim, device="cpu"):
        super(Generator, self).__init__()
        self.device = device
        self.dense = nn.Sequential(
            nn.Linear(latent_dim, 7 * 7 * 128),
            nn.ReLU(),
        )
        self.de_conv1 = nn.Sequential(
            nn.ConvTranspose2d(128, 128, (4, 4), stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
        )

        self.de_conv2 = nn.Sequential(
            nn.ConvTranspose2d(128, 128, (4, 4), stride=(2, 2), padding=(1, 1)),
            nn.BatchNorm2d(128),
            nn.ReLU(),
        )
        self.conv1 = nn.Conv2d(128, 1, (7, 7), stride=1, padding=(3, 3))
        self.tanh = nn.Tanh()

    def forward(self, x):
        # print(self.de_conv1.weight.shape)
        x = x.to(self.device)
        batch = x.shape[0]
        x = self.dense(x)
        x = x.view((batch, 128, 7, 7))
        x = self.de_conv1(x)
        x = self.de_conv2(x)
        x = self.conv1(x)
        x = self.tanh(x)
        # print(x.shape)
        return x 

